import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D  
import os
import json

def save_simulation_results(psi, u, v, omega, n_density, out_iter, data_folder='./data'):
    """
    保存模拟结果到.npz文件。

    参数:
    - psi: 二维数组，流函数。
    - u: 二维数组，x方向速度。
    - v: 二维数组，y方向速度。
    - omega: 二维数组，涡度。
    - n_density: 二维数组，密度。
    - out_iter: 整数，迭代次数，用于文件命名。
    - data_folder: 字符串，保存数据的文件夹路径，默认为'./data'。

    返回:
    - None
    """
    # 构造包含迭代次数的文件名
    file_name = f'out_iter_{out_iter}.npz'
    # 完整的文件保存路径
    full_path = os.path.join(data_folder, file_name)

    # 确保文件夹存在
    if not os.path.exists(data_folder):
        os.makedirs(data_folder)

    # 使用numpy.savez_compressed()保存所有数组到一个npz文件中
    np.savez_compressed(full_path, psi=psi, u=u, v=v, omega=omega, n_density=n_density)

def diag_uv_quiver(u,v,xx,yy,out_iter,Ra,dt):
    
    subfolder='uv_quiver' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    plt.figure(figsize=(12, 8), dpi=100)  
    # 使用quiver函数绘制矢量图  
    Q = plt.quiver(xx, yy, u, v, units='width', scale=10, color='b')  
    plt.xlabel('X')  
    plt.ylabel('Y')  
    plt.title('2D Vector Field , Re={},t={}'.format(Ra,t_diag))  
    # 定义文件名，包含idiag的值  
    filename_base = 'uv_quiver_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()
  
def diag_n_pcolor(n_density,n_init0_matrix,u,v,xx,yy,out_iter,Ra,dt):
    subfolder='n_pcolor' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    fig, ax = plt.subplots()
    n_difference=n_density-n_init0_matrix
    # 绘制伪彩色图
    plt.streamplot(xx, yy, u, v, linewidth=2)
    plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 定义文件名，包含idiag的值  
    filename_base = 'n_plot_step_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存散点图到指定路径  
    plt.savefig(full_path)  
    plt.close()

def diag_n_pcolor_filepath(n_density,u,v,xx,yy,out_iter,Ra,dt, base_folder='./n_pcolor',n_init0_matrix=None):
    """
    诊断并保存伪彩色图到指定文件夹。

    参数:
    n_density, n_init0_matrix, u, v, xx, yy -- 数据用于绘图的NumPy数组。
    out_iter -- 当前迭代步数。
    Ra, dt -- Rayleigh数和时间步长。
    base_folder -- 图片保存的基础文件夹路径。
    """
    # 目标子文件夹路径
    folder_path = os.path.join(base_folder, 'n_pcolor')
    
    # 确保子文件夹存在
    os.makedirs(folder_path, exist_ok=True)  # 使用exist_ok避免在目录已存在时抛出异常

    t_diag=out_iter*dt
    fig, ax = plt.subplots()
    if n_init0_matrix == None:
        n_difference=n_density
    else:
        n_difference=n_density-n_init0_matrix
    # 绘制流线
    plt.streamplot(xx, yy, u, v, linewidth=1)
    
    # plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), vmin=-np.abs(n_difference).max(), vmax=np.abs(n_difference).max(), shading='auto')
    # plt.colorbar(label='n',boundaries=np.linspace(-np.abs(n_difference).max(), np.abs(n_difference).max()))
    
    # 使用自动缩放的颜色条
    pcm = plt.pcolormesh(xx, yy, n_difference, cmap=plt.get_cmap('Spectral_r'), shading='auto')
    
    # 添加颜色条，不需要指定boundaries
    cbar = fig.colorbar(pcm, ax=ax, orientation='vertical')
    cbar.set_label('n')
    # 添加标题
    ax.set_title('Ra={}, t={}'.format(Ra, t_diag), fontsize=16)
    # 添加颜色条
    #cbar = fig.colorbar(cs, orientation='vertical')
    # 设置坐标轴的比例相同，使得图形保持矩形
    ax.set_aspect('equal')
    # 构建并保存文件路径
    filename = f'n_plot_step_{out_iter}.png'
    file_path = os.path.join(folder_path, filename)
    
    # 保存并关闭图表
    plt.savefig(file_path)
    plt.close()

def diag_psi_3d(psi,xx, yy,out_iter,Ra,dt):
    
    subfolder='psi_3d' #子文件夹名
    full_path_dir = os.path.join(os.getcwd(), subfolder) # 拼接当前工作目录和子文件夹路径 
    # 检查子文件夹是否存在  
    if not os.path.exists(full_path_dir):  
        # 如果不存在，则创建子文件夹  
        os.makedirs(full_path_dir)

    t_diag=out_iter*dt
    fig = plt.figure()  
    ax = fig.add_subplot(111, projection='3d')  
    
    # 绘制三维曲面  
    ax.plot_surface(xx, yy, psi, cmap='viridis', linewidth=0)  # cmap参数用于设置颜色映射  
    
    # 添加标题和坐标轴标签  
    ax.set_title('3D Plot of psi,t={:},Ra={:}'.format(t_diag,Ra))  
    ax.set_xlabel('X')  
    ax.set_ylabel('Y')  
    ax.set_zlabel('psi')  
    # 定义文件名，包含idiag的值  
    filename_base = 'psi_3d_'  
    file_extension = '.png'  
    filename = f"{filename_base}{out_iter}{file_extension}"  
  
    # 拼接完整的文件路径  
    full_path = os.path.join(os.getcwd(), subfolder, filename)  
  
    # 保存到指定路径  
    plt.savefig(full_path)  
    plt.close()

def plot_x_distribution(x_data, y_data, title, save_path, iteration=None):
    """
    绘制并保存沿x轴变化的数据分布图。
    
    参数:
    - x_data (array-like): x轴坐标数据。
    - y_data (array-like): 对应于x轴的y轴数据，例如n_density。
    - title (str): 图表的标题。
    - save_path (str): 图片保存的路径。
    - iteration (int, optional): 迭代次数或时间点，用于命名文件。默认为None，则不包含在文件名中。
    """
    # 确保保存路径的文件夹存在
    os.makedirs(os.path.dirname(save_path), exist_ok=True)
    
    # 创建图表
    plt.figure(figsize=(10, 6))
    plt.plot(x_data, y_data)
    plt.xlabel('x')
    plt.ylabel('Value')  # 这里假设y轴是某种值，您可以根据实际情况修改
    plt.title(title)
    
    filename = save_path
    
    # 保存图像
    plt.savefig(filename)
    print(f"图已保存至: {filename}")
    
    # 可选择是否显示图像
    # plt.show()
    
    # 清理，防止内存泄漏
    plt.close()

def plot_and_save_n_density_x_section(x,y, n_density, n_init0_matrix, y_slice_index, iteration, title_base="delta n_density Distribution", save_folder_base="./SOL_plot"):
    """
    绘制n_density沿x轴在Y轴中点的分布图,并保存到指定文件夹。

    参数:
    - x (numpy.ndarray): x轴坐标数组。
    -y (numpy.ndarray): y轴坐标数组。
    - n_density (numpy.ndarray): n_density的二维数组。
    - n_init0_matrix (numpy.ndarray): 初始n_density的二维数组。
    - - y_slice_index (int): 指定的y轴切片索引。
    - iteration (int): 迭代次数，用于命名保存的文件。
    - title_base (str, optional): 图标题的基底字符串，默认为"delta n_density Distribution"。
    - save_folder_base (str, optional): 保存图像的基本文件夹路径，默认为"./SOL_plot"。
    """
    # 设置默认标题和保存文件夹路径
    title = f"{title_base} along X-axis at y={y[y_slice_index]} "
    save_folder = os.path.join(save_folder_base, f"n_density_x_section_y_{y[y_slice_index]}")
    
    # 确保文件夹存在
    os.makedirs(save_folder, exist_ok=True)
    
    # 构建具体的保存路径
    #save_path = os.path.join(save_folder, f"{title.replace(' ', '_')}_{iteration}.png")
    save_path = os.path.join(save_folder, f"delta_n_density_x_section_{iteration}.png")
    
    # 计算沿Y轴中点的n_density差异
    delta_n_midpoint = n_density[y_slice_index, :] - n_init0_matrix[y_slice_index, :]
    
    # 调用函数绘制并保存分布图
    plot_x_distribution(x, delta_n_midpoint, title, save_path, iteration)
    
    print(f"沿x轴分布的图已成功绘制并保存至: {save_path}")

def plot_and_save_v_x_section(x,y, v, y_slice_index, iteration, title_base="V Distribution along X-axis", save_folder_base="./SOL_plot"):
    """
    绘制v沿x轴在Y轴中点的分布图，并保存到指定文件夹。

    参数:
    - x (numpy.ndarray): x轴坐标数组。
    -y (numpy.ndarray): y轴坐标数组。
    - v (numpy.ndarray): v的二维数组。
    - Ny (int): v数组的y轴尺寸。
    - iteration (int): 迭代次数，用于命名保存的文件。
    - title_base (str, optional): 图标题的基底字符串，默认为"V Distribution along X-axis at Y Midpoint"。
    - save_folder_base (str, optional): 保存图像的基本文件夹路径，默认为"./SOL_plot"。
    """
    # 设置标题和保存文件夹路径，直接关注v的分布
    title = f"{title_base} at y={y[y_slice_index]}"
    save_folder = os.path.join(save_folder_base, f"v_x_section_y_{y[y_slice_index]}")  # 文件夹专门用于存放v分布图
    
    # 确保文件夹存在
    os.makedirs(save_folder, exist_ok=True)
    
    # 构建保存路径，直接使用v的分布作为命名基础
    save_path = os.path.join(save_folder, f"v_x_section_{iteration}.png")
    
    # 直接取Y轴中点的v值进行绘图，无需计算差异
    v_slice = v[y_slice_index, :]
    
    # 调用函数绘制并保存分布图，确保plot_x_distribution能适应这一改变
    plot_x_distribution(x, v_slice, title, save_path, iteration)
    
    print(f"沿x轴分布的v图已成功绘制并保存至: {save_path}")

def save_to_json(data, filename):
    """
    将数据保存为JSON文件。

    参数:
    data -- 要保存的数据，应该是一个可序列化的Python字典。
    filename -- 要保存的文件名。
    """
    # 确保保存路径中的目录存在
    directory = os.path.dirname(filename)
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    # 将数据写入JSON文件
    with open(filename, 'w') as f:
        json.dump(data, f, indent=4)

def load_simulation_config(config_path='./data/SOL/simulation_config.json'):
    """
    从JSON文件加载模拟参数。

    参数:
    config_path -- 加载文件的路径。

    返回:
    config -- 包含模拟参数的字典。
    """
    # 读取JSON文件
    with open(config_path, 'r') as f:
        config = json.load(f)
    
    return config

def load_simulation_data(base_path, out_iter):
    """
    从.npz文件加载模拟数据。

    参数:
    base_path -- 保存文件的基础路径。
    out_iter -- 要加载的时间步长。

    返回:
    n, psi, omega, u, v -- 对应的模拟数据NumPy数组。
    out_iter -- 加载的数据对应的时间步长。
    """
    # 构建文件路径，包含时间步信息
    load_path = f"{base_path}/simulation_data_t_{out_iter}.npz"
    
    # 加载.npz文件
    with np.load(load_path, allow_pickle=True) as data:
        # 提取数组数据
        n = data['n']
        psi = data['psi']
        omega = data['omega']
        u = data['u']
        v = data['v']
        out_iter_loaded = data['out_iter']  # 确保加载的是正确的时间步长
    
    return n, psi, omega, u, v, out_iter_loaded

def save_initial_data(n_init0_matrix, xx, yy,x,y, save_path='./data/initial_data.npz'):
    """
    将初始条件和网格数据保存到.npz文件。

    参数:
    n_init0_matrix -- 初始密度矩阵。
    xx, yy -- 网格坐标。
    save_path -- 保存文件的路径。
    """
    # 确保保存路径中的目录存在
    directory = os.path.dirname(save_path)
    if not os.path.exists(directory):
        os.makedirs(directory)
    
    # 保存数据到.npz文件
    np.savez_compressed(save_path, n_init0_matrix=n_init0_matrix, xx=xx, yy=yy,x=x,y=y)

def load_initial_data(load_path='./data/SOL/initial_conditions.npz'):
    """
    从.npz文件加载初始条件和网格数据。

    参数:
    load_path -- 加载文件的路径。

    返回:
    n_init0_matrix -- 初始密度矩阵。
    xx, yy -- 网格坐标。
    x,y
    """
    # 加载.npz文件
    with np.load(load_path, allow_pickle=True) as data:
        # 提取数组数据
        n_init0_matrix = data['n_init0_matrix']
        xx = data['xx']
        yy = data['yy']
        x=data['x']
        y=data['y']
    
    return n_init0_matrix, xx, yy,x,y